#!/bin/bash

#
# tunfx -- by Dario Berzano <dario.berzano@gmail.com>
#
# Part of the lfscripts suite -- http://lfscripts.googecode.com/
#
# Manages SSH tunnels via a dialog interface. Works (tested) on Linux and Mac,
# but dialog is required.
#

# Program version
export VERSION="2.0.0"

# Full autossh and ssh executable path
export AUTOSSH=$(which autossh)
export SSH=$(which ssh)

# Temporary directory
export TMP="/tmp/tunfx-$(whoami)/"

# Prefix of symlinks
export PREFIX="tunfx-"

# Dialog temporary file template (for mktemp)
export DIALOG_TEMPLATE="$TMP/dialog.XXXXX"

# Default parameters for ssh and autossh
export AUTOSSH_PARAMS="-f -M \$MONPORT"
export SSH_PARAMS="-N -q"
export SSHONLY_PARAMS="-N -q -f"

# Time interval (in secs) between two consecutive network checks
export AUTOSSH_POLL=30

# Selected tunnel (name)
export TUN_SELECTED=""

# Index (ret. value of IndexFromName/IndexFromID)
export INDEX=-1

# Tunnel status (ret. value of TunnelStatus: can be On, Off)
export TUN_STATUS=""

# Tunnel command line
export TUN_THIS_CMD=""

# Settings file (it is a Bash file to source)
export SETTINGS_FILE="$HOME/.lfscripts/tunfx.conf"

# The tunnels
declare -a TUN_ID
declare -a TUN_NAME
declare -a TUN_CMD
declare -a TUN_DEP

# Tunnel counter
export TUN_COUNT=0

# Colours
export COL_YELLOW="\033[1;33m"
export COL_GREEN="\033[1;32m"
export COL_RED="\033[1;31m"
export COL_NULL="\033[m"

# Kill all ssh tunnels, disregarding settings. This is useful to kill tunnels
# that are no more in configuration files.
function KillAll() {
  WD=$(pwd)
  cd "$TMP"

  clear
  echo ""
  echo "Stopping all tunnels..."
  echo ""

  # Spaces in file name are correctly handled
  for FILE in $PREFIX*; do

    [ ! -r "$FILE" ] && break

    local TYPE=""

    if [ -L "$FILE" ]; then
      THEPID=$(ps -ax -o pid,comm | grep "$FILE" | cut -f1 -d' ')
      TYPE="symlink"
    elif [ -f "$FILE" ]; then
      THEPID=`cat "$FILE"`
      TYPE="pidfile"
    fi

    echo -n " - $FILE ($TYPE, pid=$THEPID): "
    kill -0 $THEPID 2>/dev/null
    if [ $? == 0 ]; then
      echo "stopping tunnel and removing file..."
      StopKill $THEPID
    else
      echo "not active: removing stale file..."
    fi

    rm -f "$FILE"

  done

  cd "$WD"

  echo ""
  echo "Operation completed."
  echo ""
  echo -n "Press [Enter] to continue..."
  read
}

# List available ssh tunnels
function ListTunnels() {

  # The Bash source that will store the dialog command
  DIALOG_CMD=$(mktemp $DIALOG_TEMPLATE)

  # The output file containing user's choice
  DIALOG_OUT=$(mktemp $DIALOG_TEMPLATE)

  # Heading stuff
  echo -n "dialog --no-cancel --menu" \
    "\"Welcome to tunfx v$VERSION. \nChoose a tunnel or an action:\"" \
    "0 0 0" > "$DIALOG_CMD"

  # Compose the dialog command
  for ((i=0; $i<TUN_COUNT; i++))
  do
    if [ "${TUN_NAME[$i]}" == "<separator>" ]; then
      echo -n " \"\" \"\"" >> "$DIALOG_CMD"
    else
      TunnelStatus "${TUN_NAME[$i]}"
      echo -n " \"${TUN_NAME[$i]}\" \"$TUN_STATUS\"" >> "$DIALOG_CMD"
    fi
  done

  [ $TUN_COUNT -gt 0 ] && echo -n " \"\" \"\" \"Refresh\" \"\"" >> "$DIALOG_CMD"

  # Footing stuff
  echo -n " \"Kill all tunnels\" \"\" \"Quit\" \"\"" \
    " 2> \"$DIALOG_OUT\"" >> "$DIALOG_CMD"

  # Launches the dialog
  source "$DIALOG_CMD"
  rm "$DIALOG_CMD"

  # Stores the choice into a variable
  TUN_SELECTED=$(cat "$DIALOG_OUT")
  rm $DIALOG_OUT
}

# Tunnel menu
function TunnelActions() {

  DIALOG_OUT=$(mktemp $DIALOG_TEMPLATE)

  TunnelStatus "$TUN_SELECTED"
  SSH_MODE=${TUN_MODE[$INDEX]}

  if [ "$TUN_STATUS" == "On" ]; then

    if [ "$SSH_MODE" == "autossh" ]; then
      dialog --no-cancel --menu "$TUN_SELECTED -- autossh" 0 0 0 \
        "Stop tunnel" "" \
        "Restart tunnel" "" \
        "Show command line" "" \
        "Go back" "" 2> "$DIALOG_OUT"
    else
      dialog --no-cancel --menu "$TUN_SELECTED -- ssh" 0 0 0 \
        "Stop tunnel" "" \
        "Show command line" "" \
        "Go back" "" 2> "$DIALOG_OUT"
    fi
  else
    dialog --no-cancel --menu "$TUN_SELECTED -- $SSH_MODE" 0 0 0 \
      "Start tunnel" "" \
      "Show command line" "" \
      "Go back" "" 2> "$DIALOG_OUT"
  fi

  ACTION_SELECTED=$(cat "$DIALOG_OUT")
  rm $DIALOG_OUT

  case $ACTION_SELECTED in

    "Show command line")
      IndexFromName "$TUN_SELECTED"
      dialog --msgbox "autossh $AUTOSSH_PARAMS ${TUN_CMD[$INDEX]} $SSH_PARAMS" 15 45
    ;;

    "Start tunnel")
      StartTunnel "$TUN_SELECTED"
    ;;

    "Restart tunnel")
      RestartTunnel "$TUN_SELECTED"
    ;;

    "Stop tunnel")
      StopTunnel "$TUN_SELECTED"
    ;;

    "Go back")
      # do nothing
    ;;

  esac
}

# Restarts a tunnel by sending SIGUSR1 (30) to autossh (see man autossh), or by
# stopping and restarting it if not autossh
function RestartTunnel() {

  THEPID=$(TunnelPid "$1")
  kill -30 $THEPID

  dialog --msgbox "autossh tunnel $1 restarted" 6 45
  ACTION_SELECTED="Go back"
}

# Starts a tunnel by name
function StartTunnel() {

  IndexFromName "$1"

  # Calculate dependencies
  for DEP in ${TUN_DEP[$INDEX]}
  do
    echo "scanning dep $DEP"

    # Handle multiple dependencies
    local OLD_IFS="$IFS"
    local FIRST_DEP=$(echo "$DEP" | cut -d'|' -f1)
    local DEP_SATISFIED=0
    IFS='|'
    for MULTIDEP in $DEP
    do
      IndexFromID "$MULTIDEP"
      TunnelStatus "${TUN_NAME[$INDEX]}"
      if [ "$TUN_STATUS" == "On" ]; then
        DEP_SATISFIED=1
        break
      fi
    done
    IFS="$OLD_IFS"

    # If dependency is not satisfied, automatically start the first one
    if [ "$DEP_SATISFIED" == 0 ]; then
      IndexFromID "$FIRST_DEP"
      StartTunnel "${TUN_NAME[$INDEX]}"
    fi

  done

  IndexFromName "$1"

  SSH_CMD="${TUN_CMD[$INDEX]}"

  clear
  echo -e "Starting tunnel $COL_YELLOW${TUN_NAME[$INDEX]}$COL_NULL."
  echo ""

  # See autossh documentation: AUTOSSH_* variables are a way to pass parameters
  # to it
  export AUTOSSH_PIDFILE="$TMP/$PREFIX${TUN_ID[$INDEX]}.pid"

  # Eval autossh parameters to choose an appropriate and unique monitoring port
  MONPORT=${TUN_ID[$INDEX]}
  let MONPORT=MONPORT+10000
  AUTOSSH_EVAL_PARAMS=$(eval "echo $AUTOSSH_PARAMS")

  local SSHMODE=${TUN_MODE[$INDEX]}
  local TUNCMD
  local SSH_SYMLINK="$TMP/$PREFIX${TUN_ID[$INDEX]}.ssh"

  # Which mode? In mode "autossh" pid is written in AUTOSSH_PIDFILE, while in
  # mode ssh we use the old technique of symlink...
  if [ "$SSHMODE" == "autossh" ]; then
    TUNCMD="$AUTOSSH $AUTOSSH_EVAL_PARAMS $SSH_CMD $SSH_PARAMS"
  elif [ "$SSHMODE" == "ssh" ]; then
    ln -nfs "$SSH" "$SSH_SYMLINK"
    TUNCMD="$SSH_SYMLINK $SSHONLY_PARAMS $SSH_CMD"
  else
    TUNCMD="false"
  fi

  echo "The command line is:"
  echo ""
  echo "  $TUNCMD"
  echo ""

  # Here we launch the tunnel
  $TUNCMD
  RET=$?

  echo ""

  if [ $RET != 0 ]; then
    rm -f "$SSH_SYMLINK"
    ACTION_SELECTED=""
    echo -e "${COL_RED}Error:${COL_NULL}" \
            "tunnel command returned nonzero status ($RET)."
    echo ""
    echo -n "Press [Enter] to continue..."
    read
  else
    ACTION_SELECTED="Go back"
  fi
}

# Kills gently (SIGSTOP/15), then thoroughly (SIGKILL/9) a certain PID ($1)
function StopKill() {
  local PID=$1
  local TIMEOUT=5
  local SEC
  kill -15 $PID 2> /dev/null
  for ((SEC=0; SEC<TIMEOUT; SEC++)); do
    kill -0 $PID 2> /dev/null || return
    sleep 1
  done
  kill -9 $PID 2> /dev/null
}

# Finds the specified tunnel PID by name
function TunnelPid() {

  IndexFromName "$1"
  ID=${TUN_ID[$INDEX]}

  # Try with pidfile
  THEPID=`cat "${TMP}/${PREFIX}${ID}.pid" 2> /dev/null`

  # Now try with symlink
  if [ "$THEPID" == "" ]; then
    THEPID=$(ps -ax -o pid,comm | grep "${TMP}/${PREFIX}${ID}.ssh" |
      cut -f1 -d' ')
  fi

  # Might be empty!
  echo "$THEPID"
}

# Stops selected tunnel by name
function StopTunnel() {

  IndexFromName "$1"
  ID=${TUN_ID[$INDEX]}

  THEPID=$(TunnelPid "$1")
  [ "$THEPID" != "" ] && StopKill $THEPID
  rm -f "${TMP}/${PREFIX}${ID}"*

  dialog --msgbox "Tunnel $1 stopped." 6 45
  ACTION_SELECTED="Go back"
}

# Returns a valid random port on the specified variable name
function RandomPort() {
  local R=0
  local VARNAME=$1
  while [ $R -lt 1024 ]; do
    R=$RANDOM
  done
  eval "$VARNAME=$R"
}

# Index from name
function IndexFromName() {

  # If not found it is set to -1
  INDEX=-1

  for ((i=0; $i<$TUN_COUNT; i++))
  do
    if [ "${TUN_NAME[$i]}" == "$1" ]; then
      INDEX=$i
      break
    fi
  done
}

# Index from ID
function IndexFromID() {

  # If not found it is set to -1
  INDEX=-1

  for ((i=0; $i<$TUN_COUNT; i++))
  do
    if [ "${TUN_ID[$i]}" == "$1" ]; then
      INDEX=$i
      break
    fi
  done
}

# Adds a tunnel. Called from the settings file
function AddTunnel() {
  TUN_ID[$TUN_COUNT]="$1"
  TUN_NAME[$TUN_COUNT]="$2"
  TUN_CMD[$TUN_COUNT]="$3"
  TUN_DEP[$TUN_COUNT]="$4"
  TUN_MODE[$TUN_COUNT]="ssh"
  TUN_COUNT=$((TUN_COUNT+1))
}

# Adds an autossh tunnel. Called from the settings file
function AddAutoTunnel() {
  AddTunnel "$1" "$2" "$3" "$4"
  local C=$TUN_COUNT
  let C--
  TUN_MODE[$C]="autossh"
}

# Adds a menu separator
function AddSeparator() {
  TUN_ID[$TUN_COUNT]="$1"
  TUN_NAME[$TUN_COUNT]="<separator>"
  TUN_COUNT=$((TUN_COUNT+1))
}

# Tunnel status by name, using pidfile
function TunnelStatus() {

  IndexFromName "$1"
  ID=${TUN_ID[$INDEX]}

  THEPID=$(TunnelPid "$1")
  kill -0 $THEPID 2> /dev/null

  if [ $? == 0 ]; then
    TUN_STATUS="On"
  else
    rm -f "${TMP}/${PREFIX}${ID}."*
    TUN_STATUS="Off"
  fi

  #echo "$1 --> $ID --> ($THEPID) --> $TUN_STATUS"
  #read -n1
}

# Checks for a required package and exits program if not found.
function RequiredPackage() {
  local PKG="$1"
  local EXEC="$2"

  [ "`which $EXEC 2> /dev/null`" == "" ] || return

  echo "$EXEC not found in \$PATH. It is required to run tunfx."
  echo ""
  echo "You can install it by typing (as root), for example:"
  echo ""

  if [ `uname` == "Darwin" ]; then
    echo " * Mac using fink: fink install $PKG"
    echo " * Mac using port: port install $PKG"
  else
    echo " * Debian-based:   aptitude install $PKG"
    echo " * Red-Hat based:  yum install $PKG"
  fi

  echo ""

  exit 1
}

# Check for some required packages
RequiredPackage dialog dialog
RequiredPackage autossh autossh

# Load settings
if [ -e "$SETTINGS_FILE" ]; then
  source "$SETTINGS_FILE"
else
  MSG="Settings file not found! You should create a file"
  MSG="$MSG named:\n\n${SETTINGS_FILE}\n\ncontaining all tunnels settings."
  MSG="$MSG tunfx will start with no tunnels configured."
  dialog --msgbox "$MSG" 12 45
fi

mkdir -p "$TMP"

if [ $? != 0 ]; then
  MSG="FATAL: Cannot create temporary directory $TMP. You can change it"
  MSG="$MSG to a writable location in your settings file."
  dialog --msgbox "$MSG" 7 45
  exit 1
fi

# Main loop
while [ 1 == 1 ]
do

  ListTunnels

  if [ "$TUN_SELECTED" == "Quit" ]; then

    clear
    echo "Remember that tunnels were not killed while exiting this program!"
    echo ""
    break

  elif [ "$TUN_SELECTED" == "Kill all tunnels" ]; then

    MSG="Every tunnel will be killed, also the unlisted ones. This is useful"
    MSG="$MSG to kill tunnels that are no longer in the configuration file."
    MSG="$MSG Do you want to proceed?"
    dialog --yesno "$MSG" 9 45

    if [ $? == 0 ]; then
      KillAll
    else
      dialog --msgbox "Operation cancelled." 5 45
    fi

  elif [ "$TUN_SELECTED" == "Restart opened tunnels" ]; then

    for ((i=0; $i<TUN_COUNT; i++))
    do
      TunnelStatus "${TUN_NAME[$i]}"
      if [ $TUN_STATUS == "On" ]; then
        StopTunnel "${TUN_NAME[$i]}"
        StartTunnel "${TUN_NAME[$i]}"
      fi

    done

  elif [ "$TUN_SELECTED" == "Refresh" ]; then

    # do nothing
    echo "" > /dev/null

  elif [ "$TUN_SELECTED" != "" ]; then

    ACTION_SELECTED=""

    while [ "$ACTION_SELECTED" != "Go back" ]
    do
      TunnelActions
    done

  fi

done
